import json,os,shutil
import prior
import math
from matplotlib import pyplot as plt
from expert.actions import  simple_from_shortest_path
from environment.stretch_controller import StretchController
from utils.constants.objaverse_data_dirs import OBJAVERSE_HOUSES_DIR

import matplotlib.pyplot as plt
import imageio

def build_dir(dir):
    if os.path.exists(dir):
        shutil.rmtree(dir)
    os.makedirs(dir, exist_ok=True)
    return dir

def l2_distance(a, b, ignore_y: bool = True) -> float:
    """Return the L2 distance between two points."""
    if ignore_y:
        return sum((a[k] - b[k]) ** 2 for k in ["x", "z"]) ** 0.5
    return sum((a[k] - b[k]) ** 2 for k in ["x", "y", "z"]) ** 0.5


# python -m scripts.download_objaverse_houses --save_dir /home/fd_chen/in_context_RL/spoc-robot-training/data/objaverse_houses --subset train

def load_procthor_houses(mode ="val"):
    """Load the Procthor houses dataset."""
    if mode =="samll":
        return prior.load_dataset(dataset="spoc-data", entity="spoc-robot", revision="houses-test-val")["train"]
    elif mode =="train":
        max_houses_per_split = {"train": 150000, "val": 15000, "test": 0}
        train_houses = prior.load_dataset(
            dataset="spoc-data",
            entity="spoc-robot",
            revision="local-objaverse-procthor-houses",
            path_to_splits=None,
            split_to_path={
                k: os.path.join(OBJAVERSE_HOUSES_DIR, f"{k}.jsonl.gz")
                for k in ["train", "val", "test"]
            },
            max_houses_per_split=max_houses_per_split,
        )['train']

        return train_houses
    
def save_data_json(data, save_path, should_save=True):
    """Save data to a file if should_save is True."""
    if should_save:
        try:
            with open(save_path, "w") as f:
                json.dump(data, f, indent=4)
        except IOError as e:
            print(f"Error saving data to {save_path}: {e}")


def check_agent_position(shortest_path, agent_position, threshold):
    """Check if the agent's position is within a certain threshold of the shortest path."""
    simple_points = simple_from_shortest_path(shortest_path)
    distances = [math.hypot(agent_position["position"]["x"] - point[0], agent_position["position"]["z"] - point[1]) for point in simple_points]
    return min(distances) < threshold

def save_img_from_frame(frame, save_path="./generator_data/test.jpg"):
    """Save an image from a frame."""
    plt.imsave(save_path, frame)

def positions2path(positions):
    """Convert positions to a path format."""
    return [{'x': point[0], 'y': point[1], 'z': point[2]} for point in positions]

def debug_info(DEBUG, Expert_controller: StretchController):
    """Print debug information if DEBUG is True."""
    if DEBUG:
        full_info = Expert_controller.get_current_agent_full_pose()
        info_names = ['name', 'position', 'rotation', 'cameraHorizon', 'isStanding', 'inHighFrictionArea']
        for info in info_names:
            print(f"{info}: {full_info[info]}")
        print("=" * 40)

from concurrent.futures import ThreadPoolExecutor

def save_navigation_frame(nav_frames, frame_dir, async_save=True,vedio = True):
    """Save navigation frames to a directory, optionally asynchronously."""
    
    def async_save_image(file_path, nav_frame):
        """Helper function to save an image asynchronously."""
        plt.imsave(file_path, nav_frame)
    
    if async_save:
        executor = ThreadPoolExecutor(max_workers=4)  # 设置最大线程数
        for i, nav_frame in enumerate(nav_frames):
            file_path = f"{frame_dir}/{i}.png"
            executor.submit(async_save_image, file_path, nav_frame)
    else:
        for i, nav_frame in enumerate(nav_frames):
            plt.imsave(f"{frame_dir}/{i}.png", nav_frame)


def save_as_vedio(nav_frames,save_dir):
    video_path = f"{save_dir}/navigation_video.mp4"
    imageio.mimwrite(video_path, nav_frames, fps=30)